import numpy as np
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from matplotlib import pyplot as plt
import os
import cv2
import torch

from model import cnn


def preprossed(image_path):
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

    # 使用阈值处理将图片转换为二值图像
    _, thresh = cv2.threshold(image, 128, 255, cv2.THRESH_BINARY_INV)

    # 中值滤波去噪
    median = cv2.medianBlur(thresh, 1)  # 1x1的中值滤波

    plt.imshow(median, cmap='gray')
    plt.title('median image')
    plt.axis('off')  # 不显示坐标轴
    plt.show()

    # 膨胀操作2次
    kernel = np.ones((3, 3), np.uint8)  # 3x3的结构元素
    dilated = cv2.dilate(median, kernel, iterations=2)
    # 腐蚀操作4次
    eroded = cv2.erode(dilated, kernel, iterations=4)

    plt.imshow(eroded, cmap='gray')
    plt.title('preprocessed image')
    plt.axis('off')
    plt.show()

    return eroded


img = preprossed(r"4.jpg")
def cut(img):
    #寻找轮廓
    contours, _ = cv2.findContours(img,
                        cv2.RETR_EXTERNAL,#轮廓检索模式 只检索最外层的轮廓
                        cv2.CHAIN_APPROX_SIMPLE)#轮廓近似方法 只保留轮廓的拐点信息
    #contours包含检测到的轮廓列表。每个轮廓是列表，包含轮廓上点的坐标

    #按照轮廓的x坐标排序
    contours = sorted(contours, key=lambda ctr: cv2.boundingRect(ctr)[0])
    #分割数字
    digits = []
    for contour in contours:#遍历列表中每个轮廓
        x, y, w, h = cv2.boundingRect(contour)#xy是边界框左上坐标w宽度h高度
        if w > 10 and h > 10:  # 忽略太小的噪声
            digit = img[y:y+h, x:x+w]    #x横坐标（列坐标）y纵坐标（行坐标）
            digit = cv2.resize(digit, (28, 28))  # 调整大小为28x28
            # 在digit外围加一圈10个像素的黑边
            digit = cv2.copyMakeBorder(digit,15,10,15,10, cv2.BORDER_CONSTANT, value=[0, 0, 0])
            # 裁剪回原始大小28x28
            digit = cv2.resize(digit, (28, 28))
            digit = torch.tensor(digit, dtype=torch.float32)#转为PyTorch张量
            digit = digit.unsqueeze(0).unsqueeze(0) / 255#变为[1, 1, H, W]
            digits.append(digit)    #批次为1 通道为 [0, 1] 用于前向传播

    return digits

digits = cut(img)
num_digits = len(digits)
cnn.load_state_dict(torch.load('cnn1.pkl'))
predictions = []
with torch.no_grad():
    for digit in digits:
        output = cnn(digit)
        _, predicted = torch.max(output, 1)
        predictions.append(predicted.item())
fig, axes = plt.subplots(1, num_digits, figsize=(num_digits, 2))
s=''
# 展示每个数字
for i, digit in enumerate(digits):#enumerate将一可遍历对象组合为一个索引序列
    d = digit.squeeze().numpy()  # 转换回numpy数组并去掉多余的维度

    axes[i].imshow(d, cmap='gray')
    axes[i].set_title(f'{predictions[i]}')
    s+=str(predictions[i])
    axes[i].axis('off')  # 不显示坐标轴

plt.show()
print(s)